import tensorflow as tf
import cv2 as cv
import os
import numpy as np
from tensorflow.keras.layers import Conv2D,MaxPooling2D,GlobalAveragePooling2D,Dense,Activation,Flatten,BatchNormalization
from tensorflow.keras import optimizers,metrics,losses,Sequential
from sklearn.model_selection import train_test_split

# 读取数据
def Read_Data(path):

    image_list=[]
    label=[]

    for i in os.listdir(path):
        # 拼接数据路径
        image_path=path+"\\"+i
        # print(f"image_path: {image_path}")

        image=cv.imread(image_path)
        image=cv.resize(image,(224,224))
        # 数据归一化
        image_float32=image/255.0
        image_list.append(image_float32)

        # 标签字符串
        label_str=i.split(".")[0]
        # print(f"label_str: {label_str}")

        if label_str=="cat":
            label.append(np.eye(2)[0])

        elif label_str==None:
            print("str None ......continue.......")
            continue

        else:
            label.append(np.eye(2)[1])

    return np.array(image_list),np.array(label)


class CONV_CELL(tf.keras.models.Model):
    def __init__(self, in_channels, ksize=3, padding='same'):
        super(CONV_CELL, self).__init__()
        self.convs = Sequential(
            [Conv2D(in_channels, ksize, padding=padding),
            BatchNormalization(),
            Activation("relu")]
        )

    def call(self, inputs, training=None, mask=None):
        x=self.convs(inputs)
        return x


class DarkNet_19(tf.keras.models.Model):
    def __init__(self):
        print("Initializing the darknet19 network ......")

        super(DarkNet_19, self).__init__()
        # backbone network : DarkNet-19
        # output : stride = 2, c = 32
        self.conv_1 = Sequential(
            [CONV_CELL(32),
             MaxPooling2D()]
        )

        # output : stride = 4, c = 64
        self.conv_2 = Sequential(
           [CONV_CELL(64),
            MaxPooling2D()]
        )

        # output : stride = 8, c = 128
        self.conv_3 = Sequential(
           [CONV_CELL(128),
            CONV_CELL(64),
            CONV_CELL(128),
            MaxPooling2D()]
        )

        # output : stride = 16, c = 256
        self.conv_4 = Sequential(
            [CONV_CELL(256),
             CONV_CELL(128),
             CONV_CELL(256),
             MaxPooling2D()]
        )

        # output : stride = 32, c = 512
        self.conv_5 = Sequential(
            [CONV_CELL(512),
             CONV_CELL(256),
             CONV_CELL(512),
             CONV_CELL(256),
             CONV_CELL(512),
             MaxPooling2D()]
        )

        # output : stride = 32, c = 1024
        self.conv_6 = Sequential(
            [CONV_CELL(1024),
             CONV_CELL(512),
             CONV_CELL(1024),
             CONV_CELL(512),
             CONV_CELL(1024)]
        )
        self.conv_7 = Conv2D(2,kernel_size=3)
        self.avgpool = GlobalAveragePooling2D()
        self.flatten=Flatten()

    def call(self, inputs, training=None, mask=None):
        x = self.conv_1(inputs)
        # print('conv_1 output: {}'.format(x.shape))
        x = self.conv_2(x)
        # print('conv_2 output: {}'.format(x.shape))
        x = self.conv_3(x)
        # print('conv_3 output: {}'.format(x.shape))
        x = self.conv_4(x)
        # print('conv_4 output: {}'.format(x.shape))
        x = self.conv_5(x)
        # print('conv_5 output: {}'.format(x.shape))
        x = self.conv_6(x)
        # print('conv_6 output: {}'.format(x.shape))

        x = self.conv_7(x)
        # print('conv_7 output: {}'.format(x.shape))
        x = self.avgpool(x)
        # print('avgpool output: {}'.format(x.shape))
        x=self.flatten(x)

        return x


if __name__ == '__main__':

    train_path = r'D:\DATA\CNN-DATA\cat_dog\train_test'

    image, label = Read_Data(train_path)

    print("image_shape: ", image.shape)
    print("label_shape: ", label.shape)
    print("label_data[0]: ", label[0])

    x_train,x_test,y_train,y_test=train_test_split(image,label,train_size=0.7,shuffle=True)
    print(x_train.shape)
    print(y_train.shape)
    print(x_test.shape)
    print(y_test.shape)


    model=DarkNet_19()


    model.build(input_shape=(None,224,224,3))

    model.summary()

    model.compile(
        loss=losses.binary_crossentropy,
        optimizer=optimizers.Adam(0.001),
        metrics=['acc']
    )

    #
    model.fit(x_train,y_train,epochs=2)

    score=model.evaluate(x_test,y_test)
    print("loss: ",score[0])
    print("score: ",score[1])

